Pytorch模型转Tensorflow模型的那些事

您所在的位置：网站首页 › pytorch 导入模型 › Pytorch模型转Tensorflow模型的那些事

Pytorch模型转Tensorflow模型的那些事

2023-08-19 03:13| 来源: 网络整理| 查看: 265

Pytorch模型转Tensorflow模型部署上线 pytorch模型转tensorflow流程torch模型文件转onnx文件.onnx文件转tensorflow .pb文件导入计算图，测试样例注意版本差异和环境最近在研究Query2Title模型，学术界上快速实验一般都用pytorch，但是业界部署模型上大多都还是tensorflow模型部署。也可能是自己太懒了，哈哈，pytorch用久了tensorflow有点生疏，懒的去用啦！结果就…，线上部署必须得tensorflow模型，所以就把torch的.pt文件，转成tensorflow的pb文件把，中间遇到一些问题记录下来，仅供借鉴参考！

pytorch模型转tensorflow流程

总体上pytorch转tensorflow遵循2个步骤：

torch模型文件.pt保存成onnx模型文件.onnx.onnx模型文件转换成tensorflow .pb文件 torch模型文件转onnx文件

这里直接利用torch的onnx库将.pt文件load后导出为.onnx文件即可，基本上没有什么坑，如下：

import tensorflow as tf import torch import onnx from onnx_tf.backend import prepare import os import numpy as np def torch2Onnx(): """ pytorch转onnx """ model_pytorch = YourModel() model_pytorch.load_state_dict(torch.load('xxx.pt', map_location='cpu')) # 输入placeholder dummy_input = torch.randint(0, 10000, (1, 20)) dummy_output = model_pytorch(dummy_input) print(dummy_output.shape) # Export to ONNX format torch.onnx.export(model_pytorch, dummy_input, 'model.onnx', input_names=['inputs'], output_names=['outputs'])

需要注意的是：tensorflow1.x模型是需要先定义计算图的，计算图需要输入输出节点，所以最好指定计算图的输入输出节点名称，以便后续调用tensorflow模型

.onnx文件转tensorflow .pb文件

这一步利用开源工具onnx-tf (github) 将.onnx文件转换成.pb文件。可以直接源码安装，也可以直接 pip安装，转换过程如下：

def onnx2Tensorflow(onnx_model="model.onnx", tf_model="model.tf.pb"): """ onnx转tensorflow """ """ # ---------------------报错------------------------- # Load ONNX model and convert to TensorFlow format model_onnx = onnx.load(onnx_model) tf_rep = prepare(model_onnx) # Export model as .pb file tf_rep.export_graph(tf_model) """ os.system("onnx-tf convert -i %s -o %s"%(onnx_model, tf_model))

官方教程是按照注释中的内容做的，但是我自己执行报错，但是命令行方式转换可以成功。

导入计算图，测试样例

转换成功后，我们加载计算图测试一下样例输出是否正确，如下：

def load_pb(path_to_pb): """ 加载pb文件 """ with tf.gfile.GFile(path_to_pb, 'rb') as f: # 定义计算图 graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) with tf.Graph().as_default() as graph: tf.import_graph_def(graph_def, name='') return graph def loadTFModel(): """ 加载tensorflow模型测试 """ graph = load_pb('model.tf.pb') with tf.Session(graph=graph) as sess: # Show tensor names in graph for op in graph.get_operations(): # 获取节点名称 print(op.values()) # 获取输入输出节点 output_tensor = graph.get_tensor_by_name('div_3:0') input_tensor = graph.get_tensor_by_name('inputs:0') # dummy_input = np.random.randint(0, 1000, (1, 20), dtype=np.int64) query1 = np.array([[4158, 7811, 6653,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]], dtype=np.int64) query2 = np.array([[9914, 10859, 6907, 8719, 7098, 8861, 4158, 10785, 6299, 1264, 1612, 10285, 6973, 7811, 0, 0, 0, 0, 0, 0]], dtype=np.int64) output1 = sess.run(output_tensor, feed_dict={input_tensor: query1}) output2 = sess.run(output_tensor, feed_dict={input_tensor: query2}) simi = np.dot(output1[0], output2[0]) print(simi)

很明显在测试过程中，需要给输入张量喂入输出，获取输出节点数据，这个时候我们在第一步指定的输入、输出节点变量名就派上了用场。input_tensor获取输入节点的张量值就使用了变量"inputs:0"，":0"表示该节点的第一个输出，如果该节点只有一个输出就是:0;可以看到，在第一步中我们设定的输出变量名是outputs，但是这里并不是。（个人觉得是有点bug的，所以具体的输出节点名称要从计算图上的节点中去查看。当然如果使用最新版本的onnx-tf这个问题是不存在的，由于我需要的pb文件是tf 1.12版本的，亲测1.13版本以上，没有变量名映射的bug）

注意版本差异和环境不同的tensorflow版本和onnx-tf版本需要适配：个人总结tf 1.13以上的版本可以使用最新的onnx-tf 1.7版本；tf 1.12版本可以使用onnx-tf 1.3版本。注意模型推理是在cpu还是gpu上：tensorflow模型在cpu上推理和gpu上推理存在差异！torch文件默认的输入数据格式是NxCxHxW，而TensorFlow中的函数在gpu环境上都是支持NxCxHxW格式的。但是部分函数，例如max_pool在cpu 上只支持NxHxHxC的格式。解决办法就是，如果要在cpu上推理，需要禁用gpu: os.environ[“CUDA_VISIBLE_DEVICES”]="-1" ，否则转换会默认按照gpu环境导出；当然你如果没有gpu环境，那就不存在这个问题，自动按照cpu运行环境导出。

各位看客小伙伴，有什么问题或者好的解决办法可以留言评论哦！

【本文地址】

Pytorch模型转Tensorflow模型的那些事

Pytorch模型转Tensorflow模型的那些事

今日新闻

推荐新闻